package org.seed.quantum.socket;

import org.seed.quantum.memory.BufferFactory;
import org.seed.quantum.memory.BufferPagePool;
import org.seed.quantum.nio.AsynchronousChannelProviderExt;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketOption;
import java.net.StandardSocketOptions;
import java.nio.channels.AsynchronousChannelGroup;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.nio.channels.spi.AsynchronousChannelProvider;
import java.security.InvalidParameterException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Function;



public final class SocketServer {


    /**
     * Server端服务配置。
     * <p>调用AioQuickServer的各setXX()方法，都是为了设置config的各配置项</p>
     */
    private final Config config = new Config();
    /**
     * 内存池
     */
    private BufferPagePool bufferPool;
    /**
     * 读回调事件处理
     */
    private ReadCompletionHandler aioReadCompletionHandler;
    /**
     * ConcurrentReadCompletionHandler 回调守护线程
     */
    private ThreadPoolExecutor concurrentReadCompletionHandlerExecutor;
    /**
     * 写回调事件处理
     */
    private WriteCompletionHandler aioWriteCompletionHandler;
    private BufferPagePool innerBufferPool = null;

    /**
     * 连接会话实例化Function
     */
    private Function<AsynchronousSocketChannel, SocketSession> aioSessionFunction;
    /**
     * asynchronousServerSocketChannel
     */
    private AsynchronousServerSocketChannel serverSocketChannel = null;
    /**
     * asynchronousChannelGroup
     */
    private AsynchronousChannelGroup asynchronousChannelGroup;

    private VirtualBufferFactory readBufferFactory = bufferPage -> bufferPage.allocate(config.getReadBufferSize());

    /**
     * 设置服务端启动必要参数配置
     *
     * @param port             绑定服务端口号
     * @param protocol         协议编解码
     * @param messageProcessor 消息处理器
     */
    public <T> SocketServer(int port, Protocol<T> protocol, MessageProcessor<T> messageProcessor) {
        config.setPort(port);
        config.setProtocol(protocol);
        config.setProcessor(messageProcessor);
        config.setThreadNum(Runtime.getRuntime().availableProcessors());
    }

    /**
     * @param host             绑定服务端Host地址
     * @param port             绑定服务端口号
     * @param protocol         协议编解码
     * @param messageProcessor 消息处理器
     */
    public <T> SocketServer(String host, int port, Protocol<T> protocol, MessageProcessor<T> messageProcessor) {
        this(port, protocol, messageProcessor);
        config.setHost(host);
    }

    /**
     * 启动Server端的AIO服务
     *
     * @throws IOException IO异常
     */
    public void start() throws IOException {
        start0(channel -> new SocketSession(channel, config, aioReadCompletionHandler, aioWriteCompletionHandler, bufferPool.allocateBufferPage()));
    }

    /**
     * 内部启动逻辑
     *
     * @param aioSessionFunction 实例化会话的Function
     * @throws IOException IO异常
     */
    private void start0(Function<AsynchronousSocketChannel, SocketSession> aioSessionFunction) throws IOException {
        checkAndResetConfig();
        try {
            aioWriteCompletionHandler = new WriteCompletionHandler();
            if (bufferPool == null) {
                this.bufferPool = config.getBufferFactory().create();
                this.innerBufferPool = bufferPool;
            }
            this.aioSessionFunction = aioSessionFunction;
            AsynchronousChannelProvider provider;
            if (config.isAioEnabled()) {
                aioReadCompletionHandler = new ReadCompletionHandler();
                provider = new AsynchronousChannelProviderExt();
            } else {
                concurrentReadCompletionHandlerExecutor = new ThreadPoolExecutor(1, 1,
                        49L, TimeUnit.SECONDS, new LinkedBlockingQueue<>());
                aioReadCompletionHandler = new ConcurrentReadCompletionHandler(new Semaphore(config.getThreadNum() - 1), concurrentReadCompletionHandlerExecutor);
                provider = AsynchronousChannelProvider.provider();
            }
            asynchronousChannelGroup = provider.openAsynchronousChannelGroup(config.getThreadNum(), new ThreadFactory() {
                private byte index = 0;

                @Override
                public Thread newThread(Runnable r) {
                    return bufferPool.newThread(r, "Quantum-Socket-" + (++index));
                }
            });
            this.serverSocketChannel = AsynchronousServerSocketChannel.open(asynchronousChannelGroup);
            //set socket options
            if (config.getSocketOptions() != null) {
                for (Map.Entry<SocketOption<Object>, Object> entry : config.getSocketOptions().entrySet()) {
                    this.serverSocketChannel.setOption(entry.getKey(), entry.getValue());
                }
            }
            //bind host
            if (config.getHost() != null) {
                serverSocketChannel.bind(new InetSocketAddress(config.getHost(), config.getPort()), config.getBacklog());
            } else {
                serverSocketChannel.bind(new InetSocketAddress(config.getPort()), config.getBacklog());
            }

            startAcceptThread();
        } catch (IOException e) {
            shutdown();
            throw e;
        }
        System.out.println("quantum-socket server started on port " + config.getPort() + ",threadNum:" + config.getThreadNum());
        System.out.println("quantum-socket server config is " + config);
    }

    private void startAcceptThread() {
        serverSocketChannel.accept(null, new CompletionHandler<AsynchronousSocketChannel, Void>() {
            @Override
            public void completed(AsynchronousSocketChannel channel, Void attachment) {
                try {
                    serverSocketChannel.accept(attachment, this);
                } catch (Throwable throwable) {
                    config.getProcessor().stateEvent(null, SocketStateEnum.ACCEPT_EXCEPTION, throwable);
                    failed(throwable, attachment);
                    serverSocketChannel.accept(attachment, this);
                } finally {
                    createSession(channel);
                }
            }

            @Override
            public void failed(Throwable exc, Void attachment) {
                exc.printStackTrace();
            }
        });

    }

    /**
     * 检查配置项
     */
    private void checkAndResetConfig() {
        //确保单核CPU默认初始化至少2个线程
        if (config.getThreadNum() == 1) {
            config.setThreadNum(2);
        }
    }

    /**
     * 为每个新建立的连接创建AIOSession对象
     *
     * @param channel 当前已建立连接通道
     */
    private void createSession(AsynchronousSocketChannel channel) {
        //连接成功则构造AIOSession对象
        SocketSession session = null;
        AsynchronousSocketChannel acceptChannel = channel;
        try {
            if (config.getMonitor() != null) {
                acceptChannel = config.getMonitor().shouldAccept(channel);
            }
            if (acceptChannel != null) {
                acceptChannel.setOption(StandardSocketOptions.TCP_NODELAY, true);
                session = aioSessionFunction.apply(acceptChannel);
                session.initSession(readBufferFactory.newBuffer(bufferPool.allocateBufferPage()));
            } else {
                config.getProcessor().stateEvent(null, SocketStateEnum.REJECT_ACCEPT, null);
                IOUtil.close(channel);
            }
        } catch (Exception e) {
            e.printStackTrace();
            if (session == null) {
                IOUtil.close(channel);
            } else {
                session.close();
            }
        }
    }

    /**
     * 停止服务端
     */
    public void shutdown() {
        try {
            if (serverSocketChannel != null) {
                serverSocketChannel.close();
                serverSocketChannel = null;
            }
        } catch (IOException e) {
            e.printStackTrace();
        }

        if (!asynchronousChannelGroup.isTerminated()) {
            try {
                asynchronousChannelGroup.shutdownNow();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        try {
            asynchronousChannelGroup.awaitTermination(3, TimeUnit.SECONDS);
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        if (innerBufferPool != null) {
            innerBufferPool.release();
        }
        if (concurrentReadCompletionHandlerExecutor != null) {
            concurrentReadCompletionHandlerExecutor.shutdown();
            concurrentReadCompletionHandlerExecutor = null;
        }
    }

    /**
     * 设置读缓存区大小
     *
     * @param size 单位：byte
     * @return 当前AioQuickServer对象
     */
    public final SocketServer setReadBufferSize(int size) {
        this.config.setReadBufferSize(size);
        return this;
    }

    /**
     * 是否启用 AIO 增强模式。默认：true
     *
     * @param enabled true:启用；false:禁用
     */
    public SocketServer setAioEnhance(boolean enabled) {
        config.setAioEnabled(enabled);
        return this;
    }


    /**
     * 设置Socket的TCP参数配置。
     * <p>
     * AIO客户端的有效可选范围为：<br/>
     * 2. StandardSocketOptions.SO_RCVBUF<br/>
     * 4. StandardSocketOptions.SO_REUSEADDR<br/>
     * </p>
     *
     * @param socketOption 配置项
     * @param value        配置值
     * @param <V>          配置项类型
     * @return 当前AioQuickServer对象
     */
    public final <V> SocketServer setOption(SocketOption<V> socketOption, V value) {
        config.setOption(socketOption, value);
        return this;
    }

    /**
     * 设置服务工作线程数,设置数值必须大于等于2
     *
     * @param threadNum 线程数
     * @return 当前AioQuickServer对象
     */
    public final SocketServer setThreadNum(int threadNum) {
        if (threadNum <= 1) {
            throw new InvalidParameterException("thread number must >= 2");
        }
        config.setThreadNum(threadNum);
        return this;
    }


    /**
     * 设置输出缓冲区容量
     *
     * @param bufferSize     单个内存块大小
     * @param bufferCapacity 内存块数量上限
     * @return 当前AioQuickServer对象
     */
    public final SocketServer setWriteBuffer(int bufferSize, int bufferCapacity) {
        config.setWriteBufferSize(bufferSize);
        config.setWriteBufferCapacity(bufferCapacity);
        return this;
    }

    /**
     * 设置 backlog 大小
     *
     * @param backlog backlog大小
     * @return 当前AioQuickServer对象
     */
    public final SocketServer setBacklog(int backlog) {
        config.setBacklog(backlog);
        return this;
    }

    /**
     * 设置内存池。
     * 通过该方法设置的内存池，在AioQuickServer执行shutdown时不会触发内存池的释放。
     * 该方法适用于多个AioQuickServer、AioQuickClient共享内存池的场景。
     * <b>在启用内存池的情况下会有更好的性能表现</b>
     *
     * @param bufferPool 内存池对象
     * @return 当前AioQuickServer对象
     */
    public final SocketServer setBufferPagePool(BufferPagePool bufferPool) {
        this.bufferPool = bufferPool;
        this.config.setBufferFactory(BufferFactory.IDLE_BUFFER_FACTORY);
        return this;
    }

    /**
     * 设置内存池的构造工厂。
     * 通过工厂形式生成的内存池会强绑定到当前AioQuickServer对象，
     * 在AioQuickServer执行shutdown时会释放内存池。
     * <b>在启用内存池的情况下会有更好的性能表现</b>
     *
     * @param bufferFactory 内存池工厂
     * @return 当前AioQuickServer对象
     */
    public final SocketServer setBufferFactory(BufferFactory bufferFactory) {
        this.config.setBufferFactory(bufferFactory);
        this.bufferPool = null;
        return this;
    }

    public final SocketServer setReadBufferFactory(VirtualBufferFactory readBufferFactory) {
        this.readBufferFactory = readBufferFactory;
        return this;
    }


    public static class Config {

        /**
         * 消息体缓存大小,字节
         */
        private int readBufferSize = 1024;
        /**
         * 内存块大小限制
         */
        private int writeBufferSize = 512;
        /**
         * Write缓存区容量
         */
        private int writeBufferCapacity = 64;
        /**
         * 远程服务器IP
         */
        private String host;
        /**
         * 服务器消息拦截器
         */
        private NetMonitor monitor;
        /**
         * 服务器端口号
         */
        private int port = 8888;

        /**
         * 服务端backlog
         */
        private int backlog = 1000;

        /**
         * 消息处理器
         */
        private MessageProcessor processor;
        /**
         * 协议编解码
         */
        private Protocol protocol;


        /**
         * Socket 配置
         */
        private Map<SocketOption<Object>, Object> socketOptions;

        /**
         * 线程数
         */
        private int threadNum = 1;

        /**
         * 内存池工厂
         */
        private BufferFactory bufferFactory = BufferFactory.IDLE_BUFFER_FACTORY;

        /**
         * 启用 aio 增强
         */
        private boolean aioEnabled = true;

        /**
         * 获取默认内存块大小
         *
         * @return 内存块大小
         */
        public int getWriteBufferSize() {
            return writeBufferSize;
        }

        /**
         * @param writeBufferSize 内存块大小
         */
        public void setWriteBufferSize(int writeBufferSize) {
            this.writeBufferSize = writeBufferSize;
        }

        /**
         * @return 主机地址
         */
        public String getHost() {
            return host;
        }

        /**
         * @param host 主机地址
         */
        public void setHost(String host) {
            this.host = host;
        }

        /**
         * @return 端口号
         */
        public int getPort() {
            return port;
        }

        /**
         * @param port 端口号
         */
        public void setPort(int port) {
            this.port = port;
        }

        public NetMonitor getMonitor() {
            return monitor;
        }

        public Protocol getProtocol() {
            return protocol;
        }

        public void setProtocol(Protocol protocol) {
            this.protocol = protocol;
        }

        public MessageProcessor getProcessor() {
            return processor;
        }

        /**
         * @param processor 消息处理器
         */
        public void setProcessor(MessageProcessor processor) {
            this.processor = processor;
            this.monitor = (processor instanceof NetMonitor) ? (NetMonitor) processor : null;
        }

        public int getReadBufferSize() {
            return readBufferSize;
        }

        /**
         * @param readBufferSize 读缓冲大小
         */
        public void setReadBufferSize(int readBufferSize) {
            this.readBufferSize = readBufferSize;
        }


        public Map<SocketOption<Object>, Object> getSocketOptions() {
            return socketOptions;
        }

        /**
         * @param socketOption socketOption名称
         * @param f            socketOption值
         */
        public void setOption(SocketOption socketOption, Object f) {
            if (socketOptions == null) {
                socketOptions = new HashMap<>(4);
            }
            socketOptions.put(socketOption, f);
        }

        public int getWriteBufferCapacity() {
            return writeBufferCapacity;
        }

        public void setWriteBufferCapacity(int writeBufferCapacity) {
            this.writeBufferCapacity = writeBufferCapacity;
        }

        public int getThreadNum() {
            return threadNum;
        }

        public void setThreadNum(int threadNum) {
            this.threadNum = threadNum;
        }

        public BufferFactory getBufferFactory() {
            return bufferFactory;
        }

        public void setBufferFactory(BufferFactory bufferFactory) {
            this.bufferFactory = bufferFactory;
        }

        public int getBacklog() {
            return backlog;
        }

        public void setBacklog(int backlog) {
            this.backlog = backlog;
        }

        public boolean isAioEnabled() {
            return aioEnabled;
        }

        public void setAioEnabled(boolean aioEnabled) {
            this.aioEnabled = aioEnabled;
        }

        @Override
        public String toString() {
            return "IoServerConfig{" +
                    "readBufferSize=" + readBufferSize +
                    ", writeBufferSize=" + writeBufferSize +
                    ", writeBufferCapacity=" + writeBufferCapacity +
                    ", host='" + host + '\'' +
                    ", monitor=" + monitor +
                    ", port=" + port +
                    ", backlog=" + backlog +
                    ", processor=" + processor +
                    ", protocol=" + protocol +
                    ", socketOptions=" + socketOptions +
                    ", threadNum=" + threadNum +
                    ", bufferFactory=" + bufferFactory +
                    ", aioEnhance=" + aioEnabled +
                    '}';
        }
    }
}
